package com.xwq.xwqaiagent.rag;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

/**
 * 查询翻译器
 */
@Component
public class QueryTranslation {

    private static final Logger logger = LoggerFactory.getLogger(QueryTranslation.class);

    private final TranslationQueryTransformer queryTransformer;

    public QueryTranslation(ChatModel dashscopeChatModel) {
        // 构建 ChatClient
        ChatClient.Builder chatClientBuilder = ChatClient.builder(dashscopeChatModel);

        // 自定义 Prompt 模板（可选）
        PromptTemplate customPrompt = new PromptTemplate(
                "将以下用户查询翻译成{targetLanguage}。\n" +
                        "如果查询已经是{targetLanguage}，请原样返回。\n" +
                        "不要添加任何解释或额外文本。\n" +
                        "\n原始查询: {query}\n\n翻译后的查询:\n"
        );

        // 创建并配置 TranslationQueryTransformer
        this.queryTransformer = TranslationQueryTransformer.builder()
                .chatClientBuilder(chatClientBuilder)
                .targetLanguage("中文") // 设置目标语言为中文
                .promptTemplate(customPrompt)
                .build();
    }

    /**
     * 执行查询翻译操作
     * @param prompt 原始查询文本
     * @return 翻译后的查询文本
     */
    public String doQueryTranslation(String prompt) {
        try {
            if (prompt == null || prompt.trim().isEmpty()) {
                logger.warn("输入查询为空，返回空字符串");
                return "";
            }

            Query query = new Query(prompt);
            Query transformedQuery = queryTransformer.transform(query);

            if (StringUtils.hasText(transformedQuery.text())) {
                logger.debug("翻译成功：{} -> {}", prompt, transformedQuery.text());
                return transformedQuery.text();
            } else {
                logger.warn("翻译结果为空，返回原始查询");
                return prompt;
            }

        } catch (Exception e) {
            logger.error("查询翻译失败", e);
            throw new RuntimeException("查询翻译失败", e);
        }
    }
}